{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "Start by importing TF Encrypted. We don't need to import TensorFlow as well but it's often very convenient since we can mix ordinary and encrypted computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tf_encrypted as tfe\n",
    "\n",
    "from utils import print_in_notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We only need the following step since we want to inspect the computation later in TensorBoard. It should normally be skipped to avoid the implied overhead of generating event and tracing files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tf_encrypted:Tensorflow encrypted is dumping computation traces\n",
      "INFO:tf_encrypted:Writing event and trace files to '/tmp/tensorboard'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('Tensorflow encrypted is monitoring statistics for each', 'session.run() call using a tag')\n"
     ]
    }
   ],
   "source": [
    "%load_ext tensorboard.notebook\n",
    "\n",
    "TENSORBOARD_DIR = \"/tmp/tensorboard\"\n",
    "\n",
    "tfe.set_tfe_trace_flag(True)\n",
    "tfe.set_tfe_events_flag(True)\n",
    "tfe.set_log_directory(TENSORBOARD_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computation\n",
    "\n",
    "We next define our mixed computation, in this case summing two encrypted (i.e. private) tensors coming from different input providers. Note that we are using ordinary TensorFlow to generate the inputs locally on the providers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tfe.define_private_input(\"input-provider-x\", lambda: tf.fill([2,2], 2))\n",
    "y = tfe.define_private_input(\"input-provider-y\", lambda: tf.fill([2,2], 3))\n",
    "\n",
    "z = x + y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point `z` contains the encrypted sum. To reveal only the sum to a result receiver you would normally use the following, which decrypts and executes the print function locally on the result receiver:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compute_op = tfe.define_output(\"result-receiver\", z, tf.print)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, since we are running in a notebook the above wouldn't actually display anything. To get around this we can use the `print_in_notebook` function from `utils` instead.\n",
    "\n",
    "We stress that this is only because we are running in a notebook, and using `py_func` is for instance not possible when running in an actual distributed execution context. See the [TensorFlow documentation](https://www.tensorflow.org/api_docs/python/tf/print) for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/dahl/work/tf-encrypted/tf-encrypted/examples/notebooks/utils.py:9: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "tf.py_func is deprecated in TF V2. Instead, use\n",
      "    tf.py_function, which takes a python function which manipulates tf eager\n",
      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
      "    being differentiable using a gradient tape.\n",
      "    \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/dahl/work/tf-encrypted/tf-encrypted/examples/notebooks/utils.py:9: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "tf.py_func is deprecated in TF V2. Instead, use\n",
      "    tf.py_function, which takes a python function which manipulates tf eager\n",
      "    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to\n",
      "    an ndarray (just call tensor.numpy()) but having access to eager tensors\n",
      "    means `tf.py_function`s can use accelerators such as GPUs as well as\n",
      "    being differentiable using a gradient tape.\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "compute_op = tfe.define_output(\"result-receiver\", z, print_in_notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Execution\n",
    "\n",
    "Having defined our computation we use a `tfe.Session` to run it, optionally passing in a tag when we want event and tracing files to be written.\n",
    "\n",
    "Here we first remove previous event and tracing files to make it easier to find the new runs in TensorBoard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tf_encrypted:Players: ['server0', 'server1', 'server2', 'input-provider-x', 'input-provider-y', 'result-receiver']\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[5. 5.]\n",
      " [5. 5.]]\n",
      "[[5. 5.]\n",
      " [5. 5.]]\n"
     ]
    }
   ],
   "source": [
    "!rm -rf {TENSORBOARD_DIR}\n",
    "\n",
    "with tfe.Session() as sess:\n",
    "    sess.run(compute_op, tag='sum')\n",
    "    sess.run(compute_op, tag='sum')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspection\n",
    "\n",
    "We can finally inspect our computations in TensorBoard using the tags passed to `sess.run`.\n",
    "\n",
    "Note that this is not saved in notebooks so nothing will show below unless you run this yourself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Reusing TensorBoard on port 6010 (pid 64858), started 22:39:26 ago. (Use '!kill 64858' to kill it.)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600\"\n",
       "            src=\"http://localhost:6010\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x111400550>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%tensorboard --logdir {TENSORBOARD_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
